package com.ruoyi.web.controller.chat;

import com.lzhpo.chatgpt.OpenAiUrl;
import com.lzhpo.chatgpt.entity.chat.ChatCompletionRequest;
import com.lzhpo.chatgpt.entity.chat.ChatCompletionResponse;
import com.lzhpo.chatgpt.entity.embeddings.EmbeddingRequest;
import com.lzhpo.chatgpt.entity.embeddings.EmbeddingResponse;
import com.ruoyi.common.exception.user.UserInvalidException;
import com.ruoyi.domain.RequestLog;
import com.ruoyi.service.IRequestLogService;
import com.ruoyi.web.controller.wechat.response.FromType;
import com.ruoyi.web.core.config.EncryptUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.compress.utils.Lists;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.Nullable;

import javax.servlet.http.HttpServletRequest;
import javax.validation.constraints.NotNull;
import java.util.List;

/**
 * TODO
 *
 * @author chenchuan2
 * @date 2023-01-13 11:53:07
 */
@Slf4j
@Data
public class OpenAiAuthController {

    @Autowired
    private IRequestLogService requestLogService;

    public static List<String> COMPANY = Lists.newArrayList();

    static {
        //明婵测试
//        COMPANY.add("B2E377452CFF3E04EEB9DBFDE7CF85CC");

        //初心
//        COMPANY.add("B2E377452CFF3E04781C13F752AEC808");

        //聚美
        COMPANY.add("87A61FCB7F7CF6C9D660500AC707567A");

        //吴成刚
//        COMPANY.add("A3E8F8533295313D808F0F8A60EFC1CD");
//        COMPANY.add("A0E88B431E9393A94BD387ADBF032FFC");
//        COMPANY.add("34D75C622EF33DD4DDF5AE4CF2D3991E");

        //明婵教育
//        COMPANY.add("E77E7058CE3D3C1DF5A2A3505E213D61");

        //上海有盼头网络科技有限公司
//        COMPANY.add("51B8A441FF17B394E2A91A0C2A177E3C");
//        COMPANY.add("51B8A441FF17B394E2A91A0C2A177E3D");


        //深圳前海环融联易信息科技服务有限公司，500一个月
//        COMPANY.add("5F1EDCD6E1958598F841965F04537419");

        //安光，1000，gpt4一个月
//        COMPANY.add("8DC3825E88DA8026DE87EA34AE0CD3EB");

        //测试一下
//        COMPANY.add("E19F2B59678B468741C4C68165BB4D7A");

        //苏州千万职联科技有限公司
//        COMPANY.add("03908BEB0DAA1FE60C49BC27B87E08D0");

        //刘浪
        COMPANY.add("L3908BEB0DAA1FE60C49BC27B87E08DG");

    }

    @Autowired
    protected HttpServletRequest request;

    private static final String AUTH_ID = "Authorization";
    private static final String BEARER = "Bearer ";

    private static final String FROM_ID = "from";

    protected String getOpenAiKey() {
        String authHeader = request.getHeader(AUTH_ID);
        if (StringUtils.isBlank(authHeader) || !authHeader.startsWith(BEARER)) {
            log.error("get auth is null or invalid [{}]", authHeader);
            throw new UserInvalidException();
        }
        String openAiKey = authHeader.substring(7);
        if (StringUtils.isBlank(openAiKey)) {
            log.error("get OpenAi key is null[{}]", authHeader);
            throw new UserInvalidException();
        }
        if (!COMPANY.contains(openAiKey)) {
            log.error("get OpenAi key not in company[{}]", authHeader);
            throw new UserInvalidException();
        }
        log.info("[{}] getOpenAiKey", authHeader);
        return openAiKey;
    }

    protected void saveRecord(String token, @NotNull ChatCompletionRequest chatRequest, @Nullable ChatCompletionResponse response) {
        RequestLog requestLog = new RequestLog();
        requestLog.setIntName(OpenAiUrl.CHAT_COMPLETIONS.name());
        requestLog.setToken(token);
        requestLog.setModel(chatRequest.getModel());
        requestLog.setRequest(chatRequest.toString());
        requestLog.setResponse(null == response ? "null" : response.toString());
        requestLogService.recordLog(requestLog);
    }

    protected void saveEmbeddingRecord(String token, @NotNull EmbeddingRequest embeddingRequest, @Nullable EmbeddingResponse response) {
        RequestLog requestLog = new RequestLog();
        requestLog.setIntName(OpenAiUrl.CHAT_COMPLETIONS.name());
        requestLog.setToken(token);
        requestLog.setModel(embeddingRequest.getModel());
        requestLog.setRequest(embeddingRequest.toString());
        requestLog.setResponse(null == response ? "null" : response.toString());
        requestLogService.recordLog(requestLog);
    }

    protected void saveRecordStream(String token, @NotNull ChatCompletionRequest chatRequest) {
        RequestLog requestLog = new RequestLog();
        requestLog.setIntName(OpenAiUrl.CHAT_COMPLETIONS.name());
        requestLog.setToken(token);
        requestLog.setModel(chatRequest.getModel());
        requestLog.setRequest(chatRequest.toString());
        requestLog.setResponse("stream");
        requestLogService.recordLog(requestLog);
    }
}
